from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import base64
import os
import json
   
class MessageEncryption:
    @staticmethod
    def generate_key():
        """生成16字节的随机密钥"""
        return os.urandom(16)
       
    @staticmethod
    def encrypt_message(message_dict, key):
        """使用SM4加密消息"""
        # 将消息转换为JSON字符串
        message_str = json.dumps(message_dict)
        # 创建SM4加密器
        crypt_sm4 = CryptSM4()
        crypt_sm4.set_key(key, SM4_ENCRYPT)
        # 加密消息
        encrypt_value = crypt_sm4.crypt_ecb(message_str.encode('utf-8'))
        # Base64编码便于传输
        return base64.b64encode(encrypt_value).decode('utf-8')
    
    @staticmethod
    def decrypt_message(encrypted_message, key):
           """使用SM4解密消息"""
           # Base64解码
           encrypt_value = base64.b64decode(encrypted_message)
           # 创建SM4解密器
           crypt_sm4 = CryptSM4()
           crypt_sm4.set_key(key, SM4_DECRYPT)
           # 解密消息
           decrypted_value = crypt_sm4.crypt_ecb(encrypt_value)
           # 解析JSON
           return json.loads(decrypted_value.decode('utf-8'))
       
      